# -*- coding: utf-8 -*-
"""
    @project: pythonProject
    @Author：HanYonghua
    @file： real_sinc_interpolate1.py.py
    @date：2025/7/25 15:13
    @blogs: https://www.ncatest.com.cn
"""
import numpy as np
from collections import deque
from scipy import signal

class RealTimeSincResampler:
    def __init__(self, fs_old, fs_new, L=8, window='hamming', chunk_size=1024):
        """
        实时Sinc插值重采样器
        参数:
            fs_old: 输入采样率
            fs_new: 输出采样率
            L: 插值核半长
            window: 窗函数类型
            chunk_size: 每次处理的输入块大小
        """
        self.fs_old = fs_old
        self.fs_new = fs_new
        self.L = L
        self.ratio = fs_new / fs_old
        self.chunk_size = chunk_size

        # 初始化环形缓冲区
        self.buffer = deque(maxlen=L * 2 + chunk_size + 1)
        self.buffer.extend([0.0] * (L * 2))

        # 预计算窗函数
        if window == 'hamming':
            self.win = lambda u: 0.54 + 0.46 * np.cos(np.pi * u / L)
        elif window == 'hann':
            self.win = lambda u: 0.5 + 0.5 * np.cos(np.pi * u / L)
        elif window == 'blackman':
            self.win = lambda u: 0.42 + 0.5 * np.cos(np.pi * u / L) + 0.08 * np.cos(2 * np.pi * u / L)
        else:
            self.win = lambda u: 1.0

        # 状态变量
        self.last_t = 0.0
        self.output_buffer = []
        self.output_counter = 0  # 新增：输出样本计数器
        self.total_input_samples = 0  # 新增：总输入样本数

        # 抗混叠滤波器
        if fs_new > fs_old:
            cutoff = 0.9 * (fs_old / 2)
            self.anti_alias_filter = signal.firwin(101, cutoff, fs=fs_old)
            self.filter_state = np.zeros(100)
        else:
            self.anti_alias_filter = None

    def process_chunk(self, input_chunk, is_last_chunk=False):
        """处理输入数据块，返回重采样后的数据"""
        if self.anti_alias_filter is not None:
            input_chunk, self.filter_state = signal.lfilter(
                self.anti_alias_filter, 1, input_chunk, zi=self.filter_state)

        self.buffer.extend(input_chunk)
        available_samples = len(self.buffer) - 2 * self.L

        if available_samples < 1:
            return np.array([])

        # 如果是最后一个数据块，精确计算剩余需要的输出样本
        if is_last_chunk:
            remaining_output = max(0, int(self.total_input_samples * self.ratio) - self.output_counter)
            output_samples = min(int(available_samples * self.ratio), remaining_output)
        else:
            output_samples = int(available_samples * self.ratio)

        t_new = self.last_t + np.arange(output_samples) / self.fs_new
        self.last_t = t_new[-1] + 1 / self.fs_new if output_samples > 0 else self.last_t

        output = np.zeros(output_samples)
        buffer_array = np.array(self.buffer)

        for i, t in enumerate(t_new):
            n_center = int(t * self.fs_old) + self.L
            delta = (t - (n_center - self.L) / self.fs_old) * self.fs_old
            u = np.arange(-self.L, self.L + 1) - delta
            coeffs = np.sinc(u) * self.win(u)

            n_start = max(n_center - self.L, 0)
            n_end = min(n_center + self.L + 1, len(buffer_array))

            if n_start >= n_end:
                output[i] = 0.0
                continue

            coeffs = coeffs[self.L - (n_center - n_start): self.L + (n_end - n_center)]
            output[i] = np.sum(buffer_array[n_start: n_end] * coeffs)

        consumed_samples = int(t_new[-1] * self.fs_old) - int(t_new[0] * self.fs_old) if output_samples > 0 else 0
        for _ in range(consumed_samples):
            if len(self.buffer) > 0:
                self.buffer.popleft()

        self.output_counter += len(output)
        return output

    def realtime_resample(self, input_stream):
        """
        实时处理流式输入
        使用示例:
        resampler = RealTimeSincResampler(44100, 48000)
        for chunk in audio_stream:
            output = resampler.process_chunk(chunk)
            yield output
        """
        for chunk in input_stream:
            yield self.process_chunk(chunk)


# 模拟实时音频流
def audio_stream(signal, chunk_size):
    for i in range(0, len(signal), chunk_size):
        yield signal[i:i+chunk_size]

# 使用示例
fs_old = 44100
fs_new = 48000
t = np.arange(0, 1.0, 1/fs_old)
test_signal = np.sin(2*np.pi*1000*t) + 0.5*np.sin(2*np.pi*5000*t)

resampler = RealTimeSincResampler(fs_old, fs_new, L=8, window='blackman')
resampler.total_input_samples = len(test_signal)  # 设置总输入样本数

output_signal = []
stream = audio_stream(test_signal, chunk_size=1024)
chunks = list(stream)  # 转换为列表以确定最后一个数据块

for i, chunk in enumerate(chunks):
    is_last = (i == len(chunks) - 1)
    resampled_chunk = resampler.process_chunk(chunk, is_last_chunk=is_last)
    output_signal.extend(resampled_chunk)
    print(f"处理输入块{len(chunk)}样本 -> 输出{len(resampled_chunk)}样本")

print(f"总输入: {len(test_signal)} 输出: {len(output_signal)}")
print(f"理论输出数: {int(len(test_signal)*fs_new/fs_old)}")